package com.tripcube.ai.service.impl;

import com.aliyun.bailian20231229.Client;
import com.aliyun.bailian20231229.models.*;
import com.aliyun.teautil.models.RuntimeOptions;
import com.tripcube.ai.model.dto.KnowledgeRetrievalDTO;
import com.tripcube.ai.model.vo.KnowledgeRetrievalVO;
import com.tripcube.ai.service.KnowledgeRetrievalService;
import com.tripcube.common.core.exception.ServiceException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 知识检索服务实现类
 *
 * @author tripcube
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class KnowledgeRetrievalServiceImpl implements KnowledgeRetrievalService {

    private final Client bailianClient;

    @Value("${alibaba.cloud.workspace-id}")
    private String workspaceId;

    @Override
    public KnowledgeRetrievalVO retrieveKnowledge(KnowledgeRetrievalDTO retrievalDTO) {
        try {
            // 验证参数
            validateRetrievalParams(retrievalDTO);

            // 构建检索请求
            RetrieveRequest request = new RetrieveRequest()
                    .setQuery(retrievalDTO.getQuery())
                    .setIndexId(retrievalDTO.getKnowledgeBaseId());
                    //.setTopK(retrievalDTO.getTopK() != null ? retrievalDTO.getTopK() : 5);

            RuntimeOptions runtime = new RuntimeOptions();
            long startTime = System.currentTimeMillis();

            // 调用阿里百炼API进行知识检索
            RetrieveResponse response = bailianClient.retrieveWithOptions(
                    workspaceId, request, new HashMap<>(), runtime);

            long duration = System.currentTimeMillis() - startTime;

            if (response.getBody() == null || response.getBody().getData() == null) {
                log.warn("知识检索返回空结果，查询: {}, 知识库ID: {}", retrievalDTO.getQuery(), retrievalDTO.getKnowledgeBaseId());
                return createEmptyResult(duration);
            }

            // 转换检索结果
            return convertToRetrievalVO(response.getBody().getData(), retrievalDTO.getQuery(), duration);

        } catch (Exception e) {
            log.error("知识检索失败，查询: {}, 知识库ID: {}", retrievalDTO.getQuery(), retrievalDTO.getKnowledgeBaseId(), e);
            throw new ServiceException("知识检索失败: " + e.getMessage());
        }
    }

    @Override
    public KnowledgeRetrievalVO retrieveKnowledge(String query, String knowledgeBaseId, Integer topK) {
        KnowledgeRetrievalDTO retrievalDTO = new KnowledgeRetrievalDTO();
        retrievalDTO.setQuery(query);
        retrievalDTO.setKnowledgeBaseId(knowledgeBaseId);
        retrievalDTO.setTopK(topK);
        return retrieveKnowledge(retrievalDTO);
    }

    @Override
    public KnowledgeRetrievalVO retrieveKnowledge(String query, String knowledgeBaseId, Integer topK, Double similarityThreshold) {
        // MODIFIED: 忽略threshold
        KnowledgeRetrievalDTO retrievalDTO = new KnowledgeRetrievalDTO();
        retrievalDTO.setQuery(query);
        retrievalDTO.setKnowledgeBaseId(knowledgeBaseId);
        retrievalDTO.setTopK(topK);
        return retrieveKnowledge(retrievalDTO);
    }

    @Override
    public KnowledgeRetrievalVO retrieveFromMultipleKnowledgeBases(String query, String[] knowledgeBaseIds, Integer topK) {
        List<KnowledgeRetrievalVO.RetrievalItem> allItems = new ArrayList<>();
        long totalTime = 0;

        for (String knowledgeBaseId : knowledgeBaseIds) {
            try {
                KnowledgeRetrievalDTO retrievalDTO = new KnowledgeRetrievalDTO();
                retrievalDTO.setQuery(query);
                retrievalDTO.setKnowledgeBaseId(knowledgeBaseId);
                retrievalDTO.setTopK(topK);
                KnowledgeRetrievalVO result = retrieveKnowledge(retrievalDTO);
                if (result != null && result.getItems() != null) {
                    allItems.addAll(result.getItems());
                    totalTime += result.getDuration();
                }
            } catch (Exception e) {
                log.warn("从知识库 {} 检索失败: {}", knowledgeBaseId, e.getMessage());
            }
        }

        // 按相似度分数排序并取前topK个结果
        List<KnowledgeRetrievalVO.RetrievalItem> sortedItems = allItems.stream()
                .sorted((a, b) -> Double.compare(b.getScore(), a.getScore()))
                .limit(topK != null ? topK : 5)
                .collect(Collectors.toList());

        KnowledgeRetrievalVO result = new KnowledgeRetrievalVO();
        result.setItems(sortedItems);
        result.setTotal(sortedItems.size());
        result.setDuration(totalTime);

        return result;
    }

    /**
     * 转换API响应为检索结果VO
     */
    private KnowledgeRetrievalVO convertToRetrievalVO(RetrieveResponseBody.RetrieveResponseBodyData responseData, String query, long duration) {
        KnowledgeRetrievalVO result = new KnowledgeRetrievalVO();

        try {
            List<KnowledgeRetrievalVO.RetrievalItem> items = new ArrayList<>();

            if (responseData != null && responseData.getNodes() != null) {
                for (RetrieveResponseBody.RetrieveResponseBodyDataNodes node : responseData.getNodes()) {
                    KnowledgeRetrievalVO.RetrievalItem item = convertToRetrievalItem(node);
                    if (item != null) {
                        items.add(item);
                    }
                }
            }

            result.setItems(items);
            result.setTotal(items.size());
            result.setDuration(duration);

            log.info("知识检索完成，查询: {}, 返回结果数: {}", query, items.size());

        } catch (Exception e) {
            log.error("转换检索结果失败", e);
            return createEmptyResult(duration);
        }

        return result;
    }

    /**
     * 转换单个检索项
     */
    private KnowledgeRetrievalVO.RetrievalItem convertToRetrievalItem(RetrieveResponseBody.RetrieveResponseBodyDataNodes nodeData) {
        try {
            KnowledgeRetrievalVO.RetrievalItem item = new KnowledgeRetrievalVO.RetrievalItem();

            if (nodeData != null) {
                item.setContent(nodeData.getText());
                item.setScore(nodeData.getScore() != null ? nodeData.getScore() : 0.0);
                item.setMetadata(nodeData.getMetadata());
            }

            return item;

        } catch (Exception e) {
            log.error("转换检索项失败", e);
            return null;
        }
    }

    /**
     * 创建空的检索结果
     */
    private KnowledgeRetrievalVO createEmptyResult(long duration) {
        KnowledgeRetrievalVO result = new KnowledgeRetrievalVO();
        result.setItems(new ArrayList<>());
        result.setTotal(0);
        result.setDuration(duration);
        return result;
    }

    /**
     * 验证检索参数
     */
    private void validateRetrievalParams(KnowledgeRetrievalDTO retrievalDTO) {
        if (!StringUtils.hasText(retrievalDTO.getQuery())) {
            throw new ServiceException("查询内容不能为空");
        }

        if (!StringUtils.hasText(retrievalDTO.getKnowledgeBaseId())) {
            throw new ServiceException("知识库ID不能为空");
        }

        if (retrievalDTO.getTopK() != null && retrievalDTO.getTopK() <= 0) {
            throw new ServiceException("返回结果数量必须大于0");
        }
    }
}